use crate::*;

use quote::{format_ident, quote, ToTokens};
use syn::token::{Colon, Comma, Paren};

impl From<ReturnType> for RetType {
    fn from(value: ReturnType) -> Self {
        match value {
            ReturnType::Default => Self::Empty,
            ReturnType::Type(_, t) => Self::Type(t),
        }
    }
}

impl ToTokens for RetType {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        match self {
            Self::Empty => {
                Paren::default().surround(tokens, |_| {});
            }
            Self::Type(t) => {
                t.to_tokens(tokens);
            }
        }
    }
}

enum TypedIdentListToTokensOption {
    OnlyIdent,
    OnlyType,
    Both,
}
struct ToTokensTypedIdentSlice<'a>(&'a [TypedIdent], TypedIdentListToTokensOption);

impl<'a> ToTokens for ToTokensTypedIdentSlice<'a> {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        if self.0.is_empty() {
            return;
        }

        let (rest, tail) = (&self.0[..self.0.len() - 1], self.0.last().unwrap());
        match self.1 {
            TypedIdentListToTokensOption::OnlyIdent => {
                for ti in rest {
                    ti.0.to_tokens(tokens);
                    Comma::default().to_tokens(tokens);
                }
                tail.0.to_tokens(tokens);
            }
            TypedIdentListToTokensOption::OnlyType => {
                for ti in rest {
                    ti.1.to_tokens(tokens);
                    Comma::default().to_tokens(tokens);
                }
                tail.1.to_tokens(tokens);
            }
            TypedIdentListToTokensOption::Both => {
                for ti in rest {
                    ti.0.to_tokens(tokens);
                    Colon::default().to_tokens(tokens);
                    ti.1.to_tokens(tokens);
                    Comma::default().to_tokens(tokens);
                }
                tail.0.to_tokens(tokens);
                Colon::default().to_tokens(tokens);
                tail.1.to_tokens(tokens);
            }
        }
    }
}

impl ToTokens for TypedIdent {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        self.0.to_tokens(tokens);
        Colon::default().to_tokens(tokens);
        self.1.to_tokens(tokens);
    }
}

pub(crate) fn generate(
    option: ServiceMacroOption,
    RpcServiceTrait { vis, name, fns }: RpcServiceTrait,
) -> TokenStream {
    let name_req = format_ident!("{name}Req");
    let name_resp = format_ident!("{name}Resp");

    let calls = &*(1..=fns.len())
        .map(|i| format_ident!("Call{i}"))
        .collect::<Vec<_>>();
    let args = &*fns
        .iter()
        .map(|proc| ToTokensTypedIdentSlice(&*proc.args, TypedIdentListToTokensOption::Both))
        .collect::<Vec<_>>();
    let args_names = &*fns
        .iter()
        .map(|proc| ToTokensTypedIdentSlice(&*proc.args, TypedIdentListToTokensOption::OnlyIdent))
        .collect::<Vec<_>>();
    let args_types = &*fns
        .iter()
        .map(|proc| ToTokensTypedIdentSlice(&*proc.args, TypedIdentListToTokensOption::OnlyType))
        .collect::<Vec<_>>();
    let ret_types = &*fns.iter().map(|proc| &proc.ret_type).collect::<Vec<_>>();
    let fn_names = &*fns.iter().map(|proc| &proc.fn_name).collect::<Vec<_>>();

    let mut out = TokenStream::from(quote! {
        #[derive(Debug, tweeny_rpc::serde::Serialize, tweeny_rpc::serde::Deserialize)]
        #vis enum #name_req {
            #(#calls(#args_types)),*
        }
        #[derive(Debug, tweeny_rpc::serde::Serialize, tweeny_rpc::serde::Deserialize)]
        #vis enum #name_resp {
            #(#calls(#ret_types)),*
        }
    });

    if option.server_side {
        let name_rpc_serve = format_ident!("{name}RpcServ");
        let server = TokenStream::from(quote! {
            #vis trait #name<A> {
                #(fn #fn_names(&self, from: A, #args) -> impl std::future::Future<Output = #ret_types> + Send;)*
            }
            #vis struct #name_rpc_serve<T>(#vis T);
            impl<A, T> tweeny_rpc::server::ServeRpc<A> for #name_rpc_serve<T>
            where
                A: Send + 'static,
                T: #name<A> + Sync + Send + 'static,
            {
                type Req = #name_req;
                type Resp = #name_resp;

                fn dispatch_call(
                    &self,
                    req: Self::Req,
                    from: A,
                ) -> impl std::future::Future<Output = Self::Resp> + Send {
                    async move {
                        match req {
                            #(#name_req::#calls(#args_names) => #name_resp::#calls(self.0.#fn_names(from, #args_names).await),)*
                        }
                    }
                }
            }
        });
        out.extend(server);
    }
    if option.client_side {
        let name_client = format_ident!("{name}Client");
        let name_client_handle = format_ident!("{name}ClientHandle");
        let client = TokenStream::from(quote! {
            #vis struct #name_client;
            impl #name_client {
                #vis fn new<C, SD>(
                    c: C,
                    serder: SD,
                ) -> (
                    tweeny_rpc::client::RequestingEnd<C, SD, #name_req, #name_resp>,
                    #name_client_handle,
                )
                where
                    C: tokio::io::AsyncRead + tokio::io::AsyncWrite,
                    SD: tweeny_rpc::courier::codec::ToWriterSerializer
                        + tweeny_rpc::courier::codec::FromSliceDeserializer,
                {
                    let (serving, handle) = tweeny_rpc::client::new_client(c, serder);
                    (serving, #name_client_handle(handle))
                }
            }
            #[derive(Clone)]
            #vis struct #name_client_handle(#vis tweeny_rpc::client::RequestHandle<#name_req, #name_resp>);
            impl #name_client_handle {
                #(
                    #vis async fn #fn_names(&self, #args, timeout: Option<std::time::Duration>)
                        -> Result<#ret_types, tweeny_rpc::client::error::CallError<#name_req>>
                    {
                        match self.0.call(#name_req::#calls(#args_names), timeout).await? {
                            #name_resp::#calls(resp) => Ok(resp),
                            _ => Err(tweeny_rpc::client::error::CallError::RespMismatch),
                        }
                    }
                )*
            }
        });
        out.extend(client);
    }

    out
}
